import torch.nn as nn
import torch
import numpy as np
import pickle as pk
import torchvision.utils as Util

# 此处有 bug ！
def weight_init(model):
    """
    用于初始化模型参数
    :param model: 调用这个方法进行初始化的模型
    :return: None
    """
    cname=model.__class__.__name__
    if cname.find("Conv") !=-1:
        model.weight.data.normal_(0.0,0.02)
    elif cname.find("BatcNorm") !=-1:
        model.weight.data.normal_(1.0,0.02)
        model.bias.data.fill_(0)


def batchDataTransform(batch_image:torch.Tensor,channel:int):
    """
    对批数据进行处理，包括数据归一化和维度顺序变换
    :param batch_image: 批图片 tensor
    :param channel: 通道数目
    :return:
    """
    batch_image=batch_image.float() / 255 # 数据进行标准化
    if channel==1:
        batch_image=batch_image.unsqueeze(3)

    # 将数据转换为 【N,C,H,W】 模式
    batch_image=batch_image.permute(0,3,1,2).contiguous()
    return batch_image

def getBaseGrid(imgSize:int=64,Inbatch:bool=False,batchSize:int=1,normalize=True):
    """
    get基础变形场
    :param imgSize: 图片大小
    :param Inbatch:  是否是批量数据
    :param batchSize:  批量数据的大小
    :param normalize:  是否进行归一化
    :return:
    """
    a=torch.arange(-(imgSize-1 ),imgSize,2) # 这里有bug
    if normalize:
        a=a /(imgSize-1)
    x=a.repeat(imgSize,1) #在行的维度上重复imgsize次，在列上重复 1 次
    y=x.t()

    grid=torch.cat((x.unsqueeze(0),y.unsqueeze(0)),0)

    if Inbatch:
        grid=grid.unsqueeze(0).repeat(batchSize,1,1,1)
    #grid=torch.tensor(grid,requires_grad=False)
    grid=grid.clone().detach().requires_grad_(False)

    return grid

def saveIntermediateImage(img_list:torch.Tensor, output_dir:str, n_sample:int=4,
                          id_samele=None, dim:int=-1, filename:str="myimage", nrow:int=2, normalize:bool=False):
    """
    用于保存和可视化图片
    :param img_list:  图片 tensor
    :param output_dir:  图片输出目录
    :param n_sample:  有多少个样例
    :param id_samele:  样例的下标
    :param dim:  根据维度信息来决定是否进行维度增加
    :param filename:  保存的文件名称
    :param nrow:
    :param normalize: 是否进行规范化操作
    :return:
    """
    if id_samele is None:
        images=img_list[0:n_sample,:,:,:]
    else:
        images=img_list[id_samele,:,:,:]
    if dim > 0:
        images=images[:,dim,:,:].unsqueeze(1)

    Util.save_image(images,"{}/{}.png".format(output_dir,filename),nrow=nrow,normalize=normalize,padding=2)


def readDatasetList(filedir:str):
    """
    根据数据集的地址，读取数据集文件列表
    :param filedir: 数据集文件地址
    :return: 数据集中的数据文件地址列表
    """
    with open(filedir,mode="rb") as f:
        images=pk.load(f)
    return images


### 编码器，解码器模型工具类
from autoencoder.parameter import Parameter
import torch.nn.functional as F
################ 工具类  ####################################
class Mixer(nn.Module):
    def __init__(self,parm:Parameter,in_channel:int,out_channel:int):
        """
        建立一个全连接网络
        :param parm:  超参数集合
        :param in_channel:  输入通道数目
        :param out_channel:  输出通道数目
        """
        super(Mixer, self).__init__()
        self.net=nn.Sequential(
            nn.Linear(in_channel,out_channel),
            nn.Sigmoid()
        )

    def forward(self,x:torch.Tensor):
        """
        :param x:
        :return:
        """
        out=self.net(x)
        return out

class Warper(nn.Module):
    def __init__(self,parm:Parameter):
        """
        将纹理图像进行空间变形
        :param parm:  超参数集合
        """
        super(Warper, self).__init__()
        self.parm=parm
        self.batchSize=parm.batchSize
        self.imgSize=parm.imgSize

    def forward(self,image_tensor:torch.Tensor,input_grid:torch.Tensor):
        """
        :param image_tensor: 输入图片的 tensor
        :param input_grid:  变形场
        :return:  经过变形场变换的最终图像
        """
        self.warp=input_grid.permute(0,2,3,1)
        self.output=F.grid_sample(image_tensor,self.warp)

        return self.output

class GridSpatialIntegral(nn.Module):
    def __init__(self,parm:Parameter):
        """
        变形场空间积分运算
        :param parm: 超参数集合
        """
        super(GridSpatialIntegral, self).__init__()
        self.parm=parm
        self.w=parm.imgSize

        self.filterx=torch.tensor(np.ones(shape=(1,1,1,self.w)),dtype=torch.float32,requires_grad=False)
        self.filtery=torch.tensor(np.ones(shape=(1,1,self.w,1)),dtype=torch.float32,requires_grad=False)

        if parm.useCuda:
            self.filterx=self.filterx.cuda()
            self.filtery=self.filtery.cuda()

    def forward(self, input_diffgrid:torch.Tensor):
        """
        :param input_diffgrid: 差分变形场 tensor
        :return:
        """
        x=F.conv_transpose2d(input=input_diffgrid[:,0,:,:].unsqueeze(1),weight=self.filterx,stride=1,padding=0)
        y=F.conv_transpose2d(input=input_diffgrid[:,1,:,:].unsqueeze(1),weight=self.filtery,stride=1,padding=0)
        output_grid=torch.cat((x[:,:,0:self.w,0:self.w],y[:,:,0:self.w,0:self.w]),1)

        return output_grid